package cn.edu.recommender

import com.mongodb.casbah.Imports.{MongoClient, MongoClientURI, MongoDBObject}
import org.apache.spark.sql.{Dataset, SparkSession}
import org.elasticsearch.action.admin.indices.create.CreateIndexRequest
import org.elasticsearch.action.admin.indices.delete.DeleteIndexRequest
import org.elasticsearch.action.admin.indices.exists.indices.IndicesExistsRequest
import org.elasticsearch.common.settings.Settings
import org.elasticsearch.common.transport.InetSocketTransportAddress
import org.elasticsearch.transport.client.PreBuiltTransportClient

import java.lang.Thread.sleep
import java.net.InetAddress


// 定义样例类
case class Movie(mid: Int, name: String, descri: String, timelong: String, issue: String,
                 shoot: String, language: String, genres: String, actors: String,
                 directors: String)

case class Rating(uid: Int, mid: Int, score: Double, timestamp: Int)

case class Tag(uid: Int, mid: Int, tag: String, timestamp: Int)

case class MongoConfig(uri: String, db: String)

case class ESConfig(httpHosts: String, transportHosts: String, index: String,
                    clustername: String)

object DataLoader {
  val MOVIE_DATA_PATH = "recommender/DataLoader/src/main/resources/movies.csv"
  val RATING_DATA_PATH = "recommender/DataLoader/src/main/resources/ratings.csv"
  val TAG_DATA_PATH = "recommender/DataLoader/src/main/resources/tags.csv"
  val MONGODB_MOVIE_COLLECTION = "Movie"
  val MONGODB_RATING_COLLECTION = "Rating"
  val MONGODB_TAG_COLLECTION = "Tag"
  val ES_MOVIE_INDEX = "Movie"


  def main(args: Array[String]): Unit = {
    // 定义用到的配置参数
    val config = Map(
      "spark.cores" -> "local[*]",
      "mongo.uri" -> "mongodb://localhost:27017/recommender",
      "mongo.db" -> "recommender",
      "es.httpHosts" -> "localhost:9200",
      "es.transportHosts" -> "localhost:9300",
      "es.index" -> "recommender",
      "es.cluster.name" -> "stand-alone"
    )

    val mongoConfig = MongoConfig(config("mongo.uri"), config("mongo.db"))
    val eSConfig = ESConfig(config("es.httpHosts"), config("es.transportHosts"), config("es.index"), config("es.cluster.name"))


    //新建一个到 MongoDB 的连接
    val mongoClient = MongoClient(MongoClientURI(mongoConfig.uri))
    val mongoDB = mongoClient(mongoConfig.db)
    //如果 MongoDB 中有对应的数据库，那么应该删除
    mongoDB(MONGODB_MOVIE_COLLECTION).dropCollection()
    mongoDB(MONGODB_RATING_COLLECTION).dropCollection()
    mongoDB(MONGODB_TAG_COLLECTION).dropCollection()


    // 创建一个 SparkSession
    val spark = SparkSession.builder().appName("DataLoader").master(config("spark.cores")).getOrCreate()

    // 在对 DataFrame 和 Dataset 进行操作许多操作都需要这个包进行支持
    import spark.implicits._

    // 将 Movie 数据集加载进来
    val movieRDD = spark.sparkContext.textFile(MOVIE_DATA_PATH)
    //将 MovieRDD 装换为 DataFrame
    val movieDS = movieRDD.map(item => {
      val attr = item.split("\\^")
      Movie(attr(0).toInt, attr(1).trim, attr(2).trim, attr(3).trim, attr(4).trim,
        attr(5).trim, attr(6).trim, attr(7).trim, attr(8).trim, attr(9).trim)
    }).toDS()


    // 将数据保存到 MongoDB 中
    movieDS.write
      .option("uri", mongoConfig.uri)
      .option("collection", MONGODB_MOVIE_COLLECTION)
      .mode("overwrite")
      .format("com.mongodb.spark.sql")
      .save()

    // 将 Rating 数据保存到 MongoDB 中
    val ratingRDD = spark.sparkContext.textFile(RATING_DATA_PATH)
    val ratingDS = ratingRDD.map(item => {
      val attr = item.split(",")
      Rating(attr(0).toInt, attr(1).toInt, attr(2).toDouble, attr(3).toInt)
    }).toDS()

    ratingDS
      .write
      .option("uri", mongoConfig.uri)
      .option("collection", MONGODB_RATING_COLLECTION)
      .mode("overwrite")
      .format("com.mongodb.spark.sql")
      .save()

    // 将 Tag 数据保存
    val tagRDD = spark.sparkContext.textFile(TAG_DATA_PATH)

    val tagDS = tagRDD.map(item => {
      val attr = item.split(",")
      Tag(attr(0).toInt, attr(1).toInt, attr(2).trim, attr(3).toInt)
    }).toDS()

    tagDS
      .write
      .option("uri", mongoConfig.uri)
      .option("collection", MONGODB_TAG_COLLECTION)
      .mode("overwrite")
      .format("com.mongodb.spark.sql")
      .save()

    //对数据表建索引
    mongoDB(MONGODB_MOVIE_COLLECTION).createIndex(MongoDBObject("mid" -> 1))
    mongoDB(MONGODB_RATING_COLLECTION).createIndex(MongoDBObject("uid" -> 1))
    mongoDB(MONGODB_RATING_COLLECTION).createIndex(MongoDBObject("mid" -> 1))
    mongoDB(MONGODB_TAG_COLLECTION).createIndex(MongoDBObject("uid" -> 1))
    mongoDB(MONGODB_TAG_COLLECTION).createIndex(MongoDBObject("mid" -> 1))
    //关闭 MongoDB 的连接
    mongoClient.close()


    import org.apache.spark.sql.functions._
    val groupedTagDF = tagDS.groupBy($"mid")
      .agg(concat_ws("|", collect_set($"tag"))
        .as("tags"))
      .select("mid", "tags")

    // 需要将处理后的 Tag 数据，和 Moive 数据融合，产生新的 Movie 数据
    val movieWithTagsDF = movieDS.join(groupedTagDF, Seq("mid", "mid"), "left")

    //新建一个配置
    val settings: Settings = Settings
      .builder()
      .put("cluster.name", eSConfig.clustername)
      .build()

    //新建一个 ES 的客户端
    val esClient = new PreBuiltTransportClient(settings)

    //需要将 TransportHosts 添加到 esClient 中
    val REGEX_HOST_PORT = "(.+):(\\d+)".r
    eSConfig.transportHosts.split(",").foreach {
      case REGEX_HOST_PORT(host: String, port: String) => esClient.addTransportAddress(new InetSocketTransportAddress(InetAddress.getByName(host), port.toInt))
    }


    //需要清除掉 ES 中遗留的数据
    val indicesAdminClient = esClient.admin().indices()
    if (indicesAdminClient.exists(new IndicesExistsRequest(eSConfig.index)).actionGet().isExists) {
      indicesAdminClient.delete(new DeleteIndexRequest(eSConfig.index))
    }
    indicesAdminClient.create(new CreateIndexRequest(eSConfig.index))

    //将数据写入到 ES 中
    movieWithTagsDF
      .write
      .option("es.nodes", eSConfig.httpHosts)
      .option("es.http.timeout", "100m")
      .option("es.mapping.id", "mid")
      .mode("overwrite")
      .format("org.elasticsearch.spark.sql")
      .save(eSConfig.index + "/" + ES_MOVIE_INDEX)


    sleep(300000)
    // 关闭 Spark
    spark.stop()

  }

}















